import sqlite3
import os

# 数据库路径
db_path = 'instance/database.db'

# 检查数据库是否存在
if not os.path.exists(db_path):
    print(f"数据库文件不存在: {db_path}")
    exit(1)

# 连接数据库
conn = sqlite3.connect(db_path)
cursor = conn.cursor()

try:
    # 检查 post 表中是否已有 favorites 列
    cursor.execute("PRAGMA table_info(post)")
    columns = cursor.fetchall()
    column_names = [column[1] for column in columns]
    
    if 'favorites' not in column_names:
        print("添加 favorites 列到 post 表...")
        cursor.execute("ALTER TABLE post ADD COLUMN favorites INTEGER NOT NULL DEFAULT 0")
        conn.commit()
        print("✓ 成功添加 favorites 列")
    else:
        print("✓ post 表已有 favorites 列")
    
    # 检查 post 表是否有 likes 列
    if 'likes' not in column_names:
        print("添加 likes 列到 post 表...")
        cursor.execute("ALTER TABLE post ADD COLUMN likes INTEGER NOT NULL DEFAULT 0")
        conn.commit()
        print("✓ 成功添加 likes 列")
    else:
        print("✓ post 表已有 likes 列")
    
    # 显示表的当前结构
    cursor.execute("PRAGMA table_info(post)")
    columns = cursor.fetchall()
    print("\nPost 表结构:")
    for column in columns:
        print(f"  - {column[1]} ({column[2]})")
    
    # 检查 post_favorites 表是否存在
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='post_favorites'")
    if not cursor.fetchone():
        print("\n创建 post_favorites 表...")
        cursor.execute("""
        CREATE TABLE post_favorites (
            user_id INTEGER NOT NULL,
            post_id INTEGER NOT NULL,
            PRIMARY KEY (user_id, post_id),
            FOREIGN KEY(user_id) REFERENCES user(id),
            FOREIGN KEY(post_id) REFERENCES post(id)
        )
        """)
        conn.commit()
        print("✓ 成功创建 post_favorites 表")
    else:
        print("\n✓ post_favorites 表已存在")
        cursor.execute("PRAGMA table_info(post_favorites)")
        columns = cursor.fetchall()
        print("post_favorites 表结构:")
        for column in columns:
            print(f"  - {column[1]} ({column[2]})")

except Exception as e:
    print(f"错误: {e}")
    conn.rollback()
finally:
    # 关闭连接
    conn.close()

print("\n完成!") 